﻿#pragma comment(lib, "d3d11.lib")
#pragma comment(lib, "dxgi.lib")


#include <windows.h>
#include <d3d11.h>
#include <tchar.h>
#include <sstream>
#include <chrono>
#include <iomanip>
#include <string>
#include <vector>
//引入工具CPP
#include "util.h"
#include "LogWindow.h"
#include "ProcessInjector.h"

#include "imgui.h"
#include "imgui_impl_win32.h"
#include "imgui_impl_dx11.h"

// DirectX全局变量
static ID3D11Device* g_pd3dDevice = nullptr;
static ID3D11DeviceContext* g_pd3dDeviceContext = nullptr;
static IDXGISwapChain* g_pSwapChain = nullptr;
static ID3D11RenderTargetView* g_mainRenderTargetView = nullptr;

//汇编代码
unsigned char shellcodeMe[] = {
    0xB8, 0x78, 0x55, 0x00, 0x00,                       // mov eax, 0x5578
    0x01, 0x87, 0x78, 0x55, 0x00, 0x00                  // add [edi+0x5578], eax
};





//添加日志类接口



// 创建渲染目标
void CreateRenderTarget()
{
    ID3D11Texture2D* pBackBuffer;
    g_pSwapChain->GetBuffer(0, IID_PPV_ARGS(&pBackBuffer));
    g_pd3dDevice->CreateRenderTargetView(pBackBuffer, nullptr, &g_mainRenderTargetView);
    pBackBuffer->Release();
}

// 清理渲染目标
void CleanupRenderTarget()
{
    if (g_mainRenderTargetView) {
        g_mainRenderTargetView->Release();
        g_mainRenderTargetView = nullptr;
    }
}

// 创建设备和交换链
bool CreateDeviceD3D(HWND hWnd)
{
    // 设置交换链描述
    DXGI_SWAP_CHAIN_DESC sd;
    ZeroMemory(&sd, sizeof(sd));
    sd.BufferCount = 2;
    sd.BufferDesc.Width = 0;
    sd.BufferDesc.Height = 0;
    sd.BufferDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    sd.BufferDesc.RefreshRate.Numerator = 60;
    sd.BufferDesc.RefreshRate.Denominator = 1;
    sd.Flags = DXGI_SWAP_CHAIN_FLAG_ALLOW_MODE_SWITCH;
    sd.BufferUsage = DXGI_USAGE_RENDER_TARGET_OUTPUT;
    sd.OutputWindow = hWnd;
    sd.SampleDesc.Count = 1;
    sd.SampleDesc.Quality = 0;
    sd.Windowed = TRUE;
    sd.SwapEffect = DXGI_SWAP_EFFECT_DISCARD;

    UINT createDeviceFlags = 0;
    D3D_FEATURE_LEVEL featureLevel;
    const D3D_FEATURE_LEVEL featureLevelArray[2] = {
        D3D_FEATURE_LEVEL_11_0,
        D3D_FEATURE_LEVEL_10_0
    };

    // 创建设备和交换链
    if (D3D11CreateDeviceAndSwapChain(
        nullptr,
        D3D_DRIVER_TYPE_HARDWARE,
        nullptr,
        createDeviceFlags,
        featureLevelArray,
        2,
        D3D11_SDK_VERSION,
        &sd,
        &g_pSwapChain,
        &g_pd3dDevice,
        &featureLevel,
        &g_pd3dDeviceContext) != S_OK)
    {
        return false;
    }

    CreateRenderTarget();
    return true;
}

// 清理设备
void CleanupDeviceD3D()
{
    CleanupRenderTarget();
    if (g_pSwapChain) {
        g_pSwapChain->Release();
        g_pSwapChain = nullptr;
    }
    if (g_pd3dDeviceContext) {
        g_pd3dDeviceContext->Release();
        g_pd3dDeviceContext = nullptr;
    }
    if (g_pd3dDevice) {
        g_pd3dDevice->Release();
        g_pd3dDevice = nullptr;
    }
}

// 窗口拖动相关变量
static bool isDragging = false;
static POINT dragStartPos;
static RECT dragStartWindowRect;

// Win32消息处理
extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);
LRESULT WINAPI WndProc(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam)
{
    if (ImGui_ImplWin32_WndProcHandler(hWnd, msg, wParam, lParam))
        return true;

    switch (msg)
    {
    case WM_NCHITTEST:
    {
        // 获取鼠标位置
        POINT pt = { LOWORD(lParam), HIWORD(lParam) };
        ScreenToClient(hWnd, &pt);

        // 如果鼠标在顶部30像素范围内，返回HTCAPTION让系统处理为标题栏
        if (pt.y < 30) {
            return HTCAPTION;
        }
        break;
    }

    case WM_SIZE:
        if (g_pd3dDevice != nullptr && wParam != SIZE_MINIMIZED)
        {
            CleanupRenderTarget();
            g_pSwapChain->ResizeBuffers(0, (UINT)LOWORD(lParam), (UINT)HIWORD(lParam), DXGI_FORMAT_UNKNOWN, 0);
            CreateRenderTarget();
        }
        return 0;

    case WM_SYSCOMMAND:
        if ((wParam & 0xfff0) == SC_KEYMENU) // 禁用Alt键菜单
            return 0;
        break;

    case WM_DESTROY:
        PostQuitMessage(0);
        return 0;
    }
    return DefWindowProc(hWnd, msg, wParam, lParam);
}

// 主函数
int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nShowCmd)
{

   



    // 1. 创建注入器实例
    ProcessInjector injector;



    // 创建无边框窗口
    WNDCLASSEX wc = {
        sizeof(WNDCLASSEX),
        CS_CLASSDC,
        WndProc,
        0L,
        0L,
        GetModuleHandle(nullptr),
        nullptr,
        nullptr,
        nullptr,
        nullptr,
        _T("ImGui Simple App"),
        nullptr
    };
    ::RegisterClassEx(&wc);

    // 创建窗口 - 使用无边框样式
    HWND hwnd = ::CreateWindowEx(
        WS_EX_APPWINDOW, // 让窗口出现在任务栏
        wc.lpszClassName,
        _T("Simple ImGui App"),
        WS_POPUP, // 无边框窗口
        100, 100, 800, 500, // 位置和大小
        nullptr,
        nullptr,
        wc.hInstance,
        nullptr
    );

    // 初始化Direct3D
    if (!CreateDeviceD3D(hwnd))
    {
        CleanupDeviceD3D();
        ::UnregisterClass(wc.lpszClassName, wc.hInstance);
        return 1;
    }

    // 显示窗口
    ::ShowWindow(hwnd, SW_SHOWDEFAULT);
    ::UpdateWindow(hwnd);

    // 设置ImGui上下文
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGuiIO& io = ImGui::GetIO(); (void)io;


    //设置中文显示
    io.Fonts->AddFontFromFileTTF("c:\\Windows\\Fonts\\msyh.ttc", 18.0f, nullptr,
        io.Fonts->GetGlyphRangesChineseFull());

    // 启用键盘控制
    io.ConfigFlags |= ImGuiConfigFlags_NavEnableKeyboard;

    // 设置ImGui样式
    ImGui::StyleColorsDark();

    // 设置平台/渲染器后端
    ImGui_ImplWin32_Init(hwnd);
    ImGui_ImplDX11_Init(g_pd3dDevice, g_pd3dDeviceContext);

    // 应用变量
    static int processPid;
    static int intLength;
    static char processName[128] = "";
    static char codeAddress[128] = "";
    static char dataAddress[128] = "";
    static int age = 25;
    static bool rememberMe = false;
    static std::string  messages = "";



    //创建日志窗口
    LogWindow logWindow;

    // 主循环
    bool done = false;
    while (!done)
    {
        // 处理窗口消息
        MSG msg;
        while (::PeekMessage(&msg, nullptr, 0U, 0U, PM_REMOVE))
        {
            ::TranslateMessage(&msg);
            ::DispatchMessage(&msg);
            if (msg.message == WM_QUIT)
                done = true;
        }
        if (done)
            break;

        // 开始ImGui帧
        ImGui_ImplDX11_NewFrame();
        ImGui_ImplWin32_NewFrame();
        ImGui::NewFrame();

        // 设置窗口位置和大小 - 覆盖整个客户区
        ImGui::SetNextWindowPos(ImVec2(0, 0));
        ImGui::SetNextWindowSize(ImVec2(io.DisplaySize.x, io.DisplaySize.y));

        // 开始创建ImGui窗口 - 无标题栏、无边框、不可调整大小
        ImGui::Begin("Simple Form", nullptr,
            ImGuiWindowFlags_NoTitleBar |
            ImGuiWindowFlags_NoResize |
            ImGuiWindowFlags_NoMove |
            ImGuiWindowFlags_NoCollapse
        );

        // 添加自定义标题栏
        ImGui::PushStyleColor(ImGuiCol_Button, ImVec4(0.2f, 0.2f, 0.2f, 1.0f));
        ImGui::PushStyleColor(ImGuiCol_ButtonHovered, ImVec4(0.3f, 0.3f, 0.3f, 1.0f));
        ImGui::Button("##TitleBar", ImVec2(io.DisplaySize.x, 30));
        ImGui::PopStyleColor(2);

        // 在标题栏上添加标题
        ImGui::SameLine(10);
        ImGui::Text(u8"自制版代码注入器");

        

        // 添加UI元素
        ImGui::SetCursorPosY(40);
        ImGui::Text(u8"请填写要注入的各项数据");
        ImGui::Separator();

        // 输入框
        ImGui::InputInt(u8"进程ID", &processPid);
        ImGui::InputInt(u8"手动指定指令长度", &intLength);
        ImGui::InputText(u8"进程名称", processName,IM_ARRAYSIZE(processName));
        ImGui::InputText(u8"代码地址", codeAddress, IM_ARRAYSIZE(codeAddress));
        ImGui::InputText(u8"数据地址", dataAddress, IM_ARRAYSIZE(dataAddress));

        ImGui::Separator();
  
        ImGui::SliderInt("Age", &age, 0, 100);
        ImGui::Checkbox("Remember me", &rememberMe);

        // 按钮
        if (ImGui::Button(u8"尝试注入")) {
            
                

            // 1. 创建注入器实例
            ProcessInjector injector;

            // 2. 打开目标进程 (需要提前获取PID)
            DWORD targetPid = (DWORD)(int)processPid; // 替换为目标进程的PID
            if (!injector.OpenProcess(targetPid)) {
                return 1;
            }

            // 3. 准备shellcode (示例: 简单的MessageBox shellcode)
            //std::vector<BYTE> shellcode = {
            //    0x6A, 0x00,             // push 0
            //    0x68, 0x00, 0x00, 0x00, 0x00, // push title (需要替换实际地址)
            //    0x68, 0x00, 0x00, 0x00, 0x00, // push message (需要替换实际地址)
            //    0x6A, 0x00,             // push 0
            //    0xB8, 0x00, 0x00, 0x00, 0x00, // mov eax, MessageBoxA地址
            //    0xFF, 0xD0,             // call eax
            //    0xC3                    // ret
            //    // 注意: 这是一个简化的示例，实际使用时需要正确设置字符串和函数地址
            //};
            // 
              std::vector<BYTE> shellcode = {
                    0xB8, 0x78, 0x55, 0x00, 0x00,                       // mov eax, 0x5578
                    0x01, 0x87, 0x78, 0x55, 0x00, 0x00                  // add [edi+0x5578], eax
              };


              // 4. 注入并hook
              std::wstring targetAddress = L"PlantsVsZombies.exe+0x4BA45"; // 替换为你的目标地址


              if (intLength > 0) {
              

                  if (injector.InjectAndHook(targetAddress, shellcode, intLength)) {
                      std::wcout << L"注入成功!" << std::endl;
                      logWindow.AddLogNew(u8"手动指定长度注入成功!");
                  }
                  else {
                      std::wcerr << L"注入失败!" << std::endl;
                      logWindow.AddLogNew(u8"手动指定长度注入失败!");
                  }

              }
              else {

                  if (injector.InjectAndHook(targetAddress, shellcode)) {
                      std::wcout << L"注入成功!" << std::endl;
                      logWindow.AddLogNew(u8"注入成功!");
                  }
                  else {
                      std::wcerr << L"注入失败!" << std::endl;
                      logWindow.AddLogNew(u8"注入失败!");
                  }

              }

          
           // ExeCacheHookMe((DWORD)(int)processPid,(DWORD)codeAddress, shellcodeMe, &logWindow);

            logWindow.AddLog(u8"注入结束-----------------------------");
    
        }

    

        /*logWindow.Draw(u8"日志");*/

        ImGui::SameLine();
        if (ImGui::Button(u8"退出")) {
            done = true;
        }

        // 显示点击次数
        ImGui::Text(u8"日志: ");

        ImGui::Separator();

        // 日志内容
        ImGui::BeginChild(u8"scrolling", ImVec2(0, 0), false, ImGuiWindowFlags_HorizontalScrollbar);

        for (const auto& item : logWindow.getMessage()) {
            ImGui::TextUnformatted(item.c_str());
        }

        ImGui::EndChild();

        //这是日志相关代码

        // 结束窗口
        ImGui::End();

        // 渲染
        ImGui::Render();
        const float clear_color[4] = { 0.1f, 0.1f, 0.1f, 1.0f }; // 深灰色背景
        g_pd3dDeviceContext->OMSetRenderTargets(1, &g_mainRenderTargetView, nullptr);
        g_pd3dDeviceContext->ClearRenderTargetView(g_mainRenderTargetView, clear_color);
        ImGui_ImplDX11_RenderDrawData(ImGui::GetDrawData());

        g_pSwapChain->Present(1, 0); // 启用垂直同步
    }

    // 清理
    ImGui_ImplDX11_Shutdown();
    ImGui_ImplWin32_Shutdown();
    ImGui::DestroyContext();

    CleanupDeviceD3D();
    ::DestroyWindow(hwnd);
    ::UnregisterClass(wc.lpszClassName, wc.hInstance);

    return 0;
}